# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: yuxuan

from collections import deque

import numpy as np

from coin.base.book.types import TradeSide
from coin.strategy.basis_strat.moving_average import (MovingSum,
                                                      MovingAverage,
                                                      TimeWeightedMovingAverage,
                                                      VolumeWeightedMovingAverage)
from coin.strategy.basis_strat.moving_ohlc import (MovingHigh, MovingLow, MovingOpen, MovingClose)


def feature_map():
  return {
      'BasisFeature': BasisFeature,
      'TradeImbalanceFeatures': TradeImbalanceFeatures,
      'MovingPriceFeatures': MovingPriceFeatures,
      'BookFeatures': BookFeatures,
      'TradeFeatures': TradeFeatures
  }


def get_product_params(cls, exchange_idx, products):
  if not exchange_idx:
    if cls == BasisFeature:
      exchange_idx = [0, -1]
    else:
      exchange_idx = [0]
  params = [products[e] for e in exchange_idx]
  return params


def get_features_from_config(config, parent, products):
  features = []
  for fconfig in config:
    cls = feature_map()[fconfig['name']]
    idx = fconfig.get('exchanges', None)
    product_params = get_product_params(cls, idx, products)
    params = [parent] + product_params + fconfig['params']
    features.append(cls(*params))
  return features


class Feature(object):
  def __init__(self, parent):
    self._parent = parent

  def get_mid_price(self, product):
    return self._parent.get_mid_price(product)

  def get_feed(self, product):
    return self._parent._last_feed[product]

  def get_ts(self):
    return self._parent.get_ts()

  def on_book(self, product, book):
    pass

  def on_trade(self, product, trade):
    pass

  def feature_names(self):
    raise NotImplementedError()

  def get_features(self):
    raise NotImplementedError()

  def clear(self):
    pass

  def on_sampled(self):
    features = self.get_features()
    self.clear()
    return features


class BasisFeature(Feature):
  def __init__(self, parent, ref_product, trade_product, time_windows):
    Feature.__init__(self, parent)
    self._ref_product = ref_product
    self._trade_product = trade_product
    self._time_windows = time_windows
    self._ma = []
    for time_window in self._time_windows:
      self._ma.append(MovingAverage(time_window))

  def _get_basis(self):
    rp = self.get_mid_price(self._ref_product)
    tp = self.get_mid_price(self._trade_product)
    return rp / tp

  def on_book(self, product, book):
    if not self._parent.ready():
      return
    if product in [self._ref_product, self._trade_product]:
      basis = self._get_basis()
      for ma in self._ma:
        ma.update(self.get_ts(), basis)

  def feature_names(self):
    names = []
    for i, tw in enumerate(self._time_windows):
      names.append('norm_basis_step_%d' % i)
      names.append('norm_basis_std_%d' % i)
    return names

  def get_features(self):
    basis = self._get_basis()
    res = []
    for ma in self._ma:
      if ma.size() == 0:
        res.append(np.nan)
      else:
        mean, std = ma.get_statistics()
        if std > 0:
          res.append(basis / mean - 1)
        else:
          res.append(np.nan)
        res.append(std / mean)
    return res


class TradeImbalanceFeatures(Feature):
  def __init__(self, parent, product, time_windows):
    Feature.__init__(self, parent)
    self._product = product
    self._time_windows = time_windows
    self._buy_ms = [MovingSum(tw) for tw in self._time_windows]
    self._sell_ms = [MovingSum(tw) for tw in self._time_windows]

  def on_trade(self, product, trade):
    if product == self._product:
      if trade.side == TradeSide.TRADE_BUY_SIDE:
        for ma in self._buy_ms:
          ma.update(self.get_ts(), trade.qty)
      else:
        for ma in self._sell_ms:
          ma.update(self.get_ts(), trade.qty)

  def feature_names(self):
    names = []
    for i, tw in enumerate(self._time_windows):
      names.append('trade_imb_%s' % i)
    return names

  def get_features(self):
    features = []
    for bms, sms in zip(self._buy_ms, self._sell_ms):
      bs = bms.get_statistics()
      ss = sms.get_statistics()
      features.append((bs - ss) / (bs + ss + 1e-5))
    return features


class TradeFeatures(Feature):
  def __init__(self, parent, product, time_window):
    Feature.__init__(self, parent)
    self._product = product
    self._prices = [
        MovingHigh(time_window),
        MovingLow(time_window),
        MovingOpen(time_window),
        MovingClose(time_window)
    ]
    self._vwap = VolumeWeightedMovingAverage(time_window)

  def on_trade(self, product, trade):
    if product == self._product:
      for p in self._prices:
        p.update(self.get_ts(), trade.price)
      self._vwap.update(self.get_ts(), (trade.price, trade.qty))

  def feature_names(self):
    return ['high', 'low', 'open', 'close', 'vwap']

  def get_features(self):
    mid = self.get_mid_price(self._product)
    features = []
    for p in self._prices:
      p.update(self.get_ts())
      features.append((p.get() or mid) / mid - 1.0)
    vwap = np.nan_to_num(self._vwap.get_statistics(self.get_ts()) / mid - 1.0)
    features.append(vwap)
    return features


class MovingPriceFeatures(Feature):
  def __init__(self, parent, product, time_windows):
    Feature.__init__(self, parent)
    self._product = product
    self._time_windows = time_windows
    self._ma = [TimeWeightedMovingAverage(tw) for tw in self._time_windows]

  def on_book(self, product, book):
    if product == self._product:
      price = self.get_mid_price(self._product)
      for ma in self._ma:
        ma.update(self.get_ts(), price)

  def feature_names(self):
    names = []
    for i, tw in enumerate(self._time_windows):
      names.append('spot_ma_change_%d' % i)
    return names

  def get_features(self):
    price = self.get_mid_price(self._product)
    features = []
    for ma in self._ma:
      mean = ma.get_statistics(self.get_ts())
      features.append(price / mean - 1.0)
    return features


class BookFeatures(Feature):
  def __init__(self, parent, product):
    Feature.__init__(self, parent)
    self._product = product
    self._feeds = deque(maxlen=2)
    self._data = MovingAverage(10 * 60 * 1e9)

  def on_book(self, product, book):
    if product == self._product:
      bids, asks = self.get_feed(product)
      self._feeds.append((bids, asks))

  def feature_names(self):
    return ['qty_taken', 'bqp', 'aqp', 'book_imb']

  def get_features(self):
    qty = 0
    if len(self._feeds) == 2:
      bids0, asks0 = self._feeds[0]
      bids1, asks1 = self._feeds[1]
      for p, q in asks0:
        if p is None or p >= asks1[0][0]:
          break
        qty += q
      for p, q in bids0:
        if p is None or p <= bids1[0][0]:
          break
        qty -= q
    if qty != 0:
      self._data.update(self.get_ts(), qty)
    if self._data.size() > 100:
      _, std = self._data.get_statistics()
      if std != 0:
        zs = qty / std
      else:
        zs = 0
    else:
      zs = 0

    if len(self._feeds) > 0:
      bids, asks = self._feeds[-1]
      pq = 0
      bq = 0
      for bid in bids[:3]:
        if bid[0] is None:
          break
        pq += bid[0] * bid[1]
        bq += bid[1]
      bqp = pq / bq
      pq = 0
      aq = 0
      for ask in asks[:3]:
        if ask[0] is None:
          break
        pq += ask[0] * ask[1]
        aq += ask[1]
      aqp = pq / aq

      return [zs, np.log(bids[0][0] / bqp), np.log(aqp / asks[0][0]), (bq - aq) / (bq + aq)]
    else:
      return [np.nan] * len(self.feature_names())
